package com.junjie.checkin.websocket.handler;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.security.Principal;
import java.util.Map;

/**
 * @author 洪浚杰
 * @description
 * @date 2025/6/24 21:53
 */
@Component
public class CustomPrincipalHandshakeHandler extends DefaultHandshakeHandler {
    @Override
    protected Principal determineUser(ServerHttpRequest request,
                                      WebSocketHandler wsHandler,
                                      Map<String, Object> attributes) {
        // 1. 从 WebSocket 属性获取
        Principal principal = (Principal) attributes.get("user");
        if (principal != null) return principal;

        // 2. 从 HTTP 请求属性获取（兼容 STOMP 层）
        if (request instanceof ServletServerHttpRequest) {
            principal = (Principal) ((ServletServerHttpRequest) request)
                    .getServletRequest().getAttribute("user");
        }
        return principal != null ? principal : super.determineUser(request, wsHandler, attributes);
    }
}